import pika
import logging
import time
import json
import requests
import threading
 
# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s:%(message)s')
 
# RabbitMQ连接配置
rabbitmq_url = "amqp://user:password@172.16.168.58:5672/"
queue_name = "imgurl"
queue_name1 = "imgdesc"
 
# API配置
api_endpoint = "http://172.16.168.58/v1/chat/completions"
model_path = "Qwen/Qwen2.5-VL-7B-Instruct"
base_url = "http://fileserver/"
query_text = "描述这张图片的内容，如果有人类，描述人类的动作、肤色、民族、神情、衣着等要素"
 
# 工作通道
job_channel = None
 
def connect_to_rabbitmq():
   params = pika.URLParameters(rabbitmq_url)
   connection = pika.BlockingConnection(params)
   return connection
 
def callback(ch, method, properties, body):
   logging.info(f"Received a message: {body.decode()}")
   image_url = base_url + body.decode()
 
   # 调用API进行分析
   start_time = time.time()
   content, usage, err = query_image(image_url, query_text)
   end_time = time.time()
   total_elapsed_ms = int((end_time - start_time) * 1000)
   logging.info(f"Sent message to {queue_name1} with total elapsed time: {total_elapsed_ms} ms")
   if err:
       logging.error(f"API request failed: {err}")
   else:
       # 将结果发送到另一个队列
       job_channel.basic_publish(
           exchange='',
           routing_key=queue_name1,
           body=json.dumps({
               'imgurl': image_url,
               'content': content,
               'usage': usage,
               'total_elapsed_ms': total_elapsed_ms
           },ensure_ascii=False).encode('utf-8')
       )
       logging.info(f"Sent message to {queue_name1}      {body}")
 
   # 手动确认消息
   ch.basic_ack(delivery_tag=method.delivery_tag)
 
def query_image(image_url, query_text):
   # 构造请求体
   req_body = {
       "model": model_path,
       "messages": [
           {"role": "system", "content": "You are a helpful assistant."},
           {"role": "user", "content": [{"type": "image_url", "image_url": {"url": image_url}}, {"type": "text", "text": query_text}]}
       ]
   }
 
   # 发送请求
   try:
       response = requests.post(api_endpoint, json=req_body)
       response.raise_for_status()
   except requests.RequestException as e:
       return None, None, e
 
   # 解析响应
   try:
       api_resp = response.json()
       if len(api_resp['choices']) == 0:
           raise ValueError("API returned empty choices array")
       content = api_resp['choices'][0]['message']['content']
       usage = api_resp['usage']
       return content, usage, None
   except (KeyError, ValueError) as e:
       return None, None, e
 
def consume_messages():
   connection = connect_to_rabbitmq()
   channel = connection.channel()
 
   # 声明队列
   channel.queue_declare(queue=queue_name, durable=True)
 
   # 设置QoS
   channel.basic_qos(prefetch_count=1)
 
   # 消费消息
   channel.basic_consume(queue=queue_name, on_message_callback=callback, auto_ack=False)
 
   logging.info(" [*] Waiting for messages. To exit press CTRL+C")
   channel.start_consuming()
 
def produce_messages():
   connection = connect_to_rabbitmq()
   global job_channel
   job_channel = connection.channel()
 
   # 声明队列
   job_channel.queue_declare(queue=queue_name1, durable=True)
 
   # 等待消费者处理消息
   while True:
       time.sleep(1)
 
if __name__ == "__main__":
   # 启动消费者和生产者线程
   consumer_thread = threading.Thread(target=consume_messages)
   producer_thread = threading.Thread(target=produce_messages)
   consumer_thread.start()
   producer_thread.start()
 
   # 等待线程结束
   consumer_thread.join()
   producer_thread.join()
